import os

import pandas as pd


file_config = {
    "股票市场交易-日个股回报率文件.csv": {
        "output": "stock_daily.pkl",
        "converters": None,
        "rename": {
            "证券代码": "stkcd",
            "交易日期": "date",
            "考虑现金红利再投资的日个股回报率": "returns",
        },
        "keeping": ["date", "stkcd", "returns"],
        "start": "2010-01-01",
    }
}

cur_dir = os.path.dirname(__file__)
all_files = os.listdir(cur_dir)
for file in all_files:
    if file.endswith(".csv"):
        data = pd.read_csv(
            os.path.join(cur_dir, file), converters=file_config[file]["converters"]
        )
        if "rename" in file_config[file].keys():
            data.rename(columns=file_config[file]["rename"], inplace=True)
        if "keeping" in file_config[file].keys():
            data = data[file_config[file]["keeping"]]
        if "start" in file_config[file].keys():
            data = data[data["date"] >= file_config[file]["start"]]
        data.to_pickle(os.path.join(cur_dir, file_config[file]["output"]), "gzip")
